import cv2
import rclpy
from cv_bridge import CvBridge
from rclpy.node import Node
from sensor_msgs.msg import Image


# ========== 图像订阅节点 ==========
class ImageSubscriber(Node):
    def __init__(self):
        super().__init__('image_subscriber')
        # 创建订阅者
        self.subscription = self.create_subscription(
            Image,
            'virtual_image',
            self.image_callback,
            10  # 队列大小
        )
        self.bridge = CvBridge()
        self.get_logger().info('图像订阅节点已启动')
        self.window_name = "ROS2图像显示"

        # cv2.namedWindow(self.window_name, cv2.WINDOW_AUTOSIZE)

    def image_callback(self, msg):
        """ 接收图像并显示 """
        try:
            # 转换ROS消息为OpenCV格式（BGR8编码）
            cv_image = self.bridge.imgmsg_to_cv2(msg, "bgr8")
            # 在图像左上角添加时间戳
            text = f"Time: {msg.header.stamp.sec}.{msg.header.stamp.nanosec}"
            cv2.putText(cv_image, text, (10, 30),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
            cv2.imwrite("swd.jpg", cv_image)
            # 显示图像
            # cv2.imshow(self.window_name, cv_image)
            # cv2.waitKey(1)  # 必要！用于刷新OpenCV窗口

        # 捕获cv_bridge转换异常
        except Exception as e:
            self.get_logger().error(f'图像转换失败: {str(e)}')


# ========== 主函数 ==========
def main(args=None):
    rclpy.init(args=args)

    # 启动订阅节点（独立线程）
    subscriber = ImageSubscriber()
    executor = rclpy.executors.MultiThreadedExecutor()
    executor.add_node(subscriber)

    try:
        executor.spin()  # 多线程执行
    except KeyboardInterrupt:
        pass
    finally:
        # 清理资源
        # cv2.destroyAllWindows()
        subscriber.destroy_node()
        rclpy.shutdown()
